import Router from 'koa-router';
import {getCollection} from '../chroma/index.js'
const knowledge_base_search = new Router();

knowledge_base_search.get('/', async(ctx) => {
    const { knowledge_base, query } = ctx.query
    let docs = null
    try{
        const vectorStore = getCollection(knowledge_base)
        if(!!vectorStore){
            docs = await vectorStore.similaritySearchWithScore(query,5)
            let tmp = []
            docs.forEach(doc => {
                tmp.push({
                    text: doc[0].pageContent,
                    source: doc[0].metadata.source.split('/').pop(),
                    score: doc[1]
                })
            })
            ctx.status = 200
            ctx.body = tmp
        }else {
            ctx.status = 500
            ctx.body = {
                code: 500,
                message: '知识库不存在！'
            }
        }
    }catch (e){
        ctx.status = 503
        ctx.body = {
            code: 503,
            message: '模型未上线！'
        }
        console.log(e.name + "：" + e.message)
    }
})
export default knowledge_base_search;
